package com.akita.spring.ai.controller;

import com.akita.spring.ai.enums.BizTypeEnum;
import com.akita.spring.ai.repository.ChatHistoryRepository;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.model.Media;
import org.springframework.util.MimeType;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Objects;

@Slf4j
@RestController
@RequestMapping("/ai")
public class ChatController {

    @Resource
    private ChatClient ollamaChatClient;

    @Resource
    private ChatClient aliMultiModalChatClient;

    @Resource
    private ChatClient aliyunChatClient;

    @Resource
    private ChatHistoryRepository chatHistoryRepository;

    @RequestMapping(value = "/chat", produces = "text/html;charset=UTF-8")
    public Flux<String> chat(@RequestParam("prompt") String prompt,
                             @RequestParam("chatId") String chatId,
                             @RequestParam(value = "files", required = false) List<MultipartFile> files) {
        // 保存聊天记录
        chatHistoryRepository.save(BizTypeEnum.CHAT.getType(), chatId);

        if (CollectionUtils.isEmpty(files)) {
            // 发送纯文本消息
            log.info("调用本地ollama模型，prompt：{}", prompt);
            return textChat(prompt, chatId);  // 调用本地ollama模型
        } else {
            log.info("调用远程大模型，prompt：{}", prompt);
            return multiModalChat(prompt, chatId, files);   // 涉及到文本处理调用远程大模型
        }
    }

    /**
     * 处理文本聊天
     *
     * @param prompt 聊天内容
     * @param chatId 当前会话组的ID
     * @return 响应内容的Flux流
     */
    private Flux<String> textChat(String prompt, String chatId) {
        return aliMultiModalChatClient.prompt()
                .user(prompt)
                .advisors(advisorSpec -> advisorSpec.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }

    /**
     * 处理多媒体聊天
     *
     * @param prompt 聊天内容
     * @param chatId 当前会话组的ID
     * @param files  上传的多媒体文件列表
     * @return 响应内容的Flux流
     */
    private Flux<String> multiModalChat(String prompt, String chatId,
                                        List<MultipartFile> files) {
        // 解析多媒体文件
        List<Media> mediaList = files.stream()
                .map(file -> new Media(
                        MimeType.valueOf(Objects.requireNonNull(file.getContentType())),
                        file.getResource())
                )
                .toList();

        return aliMultiModalChatClient.prompt()
                .user(p -> p.text(prompt).media(mediaList.toArray(Media[]::new)))
                .advisors(advisorSpec -> advisorSpec.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }
}
